# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     train
   Description :   
   Author :       lth
   date：          2022/12/13
-------------------------------------------------
   Change Activity:
                   2022/12/13 17:17: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import numpy as np
import torch
from PIL import Image
from torch import optim
from torch.backends import cudnn
from torch.utils.data import DataLoader
from torchvision import transforms
from tqdm import tqdm

from confg import GetConfig
from datalist import ACEDataset
from model import ResnetEncoderDecoder
from utils import ACELoss

total_loss = 999999999999999999999999999

mem = "-abcdefghijklmnopqrstuvwxyz"


class Train:
    def __init__(self):
        self.args = GetConfig()
        print(f"-----------{self.args.project_name}-------------")
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if use_cuda else "cpu")

        # data
        kwargs = {"num_workers": 8, "pin_memory": True} if use_cuda else {"num_workers": 4, "pin_memory": False}

        self.train_dataloader = DataLoader(ACEDataset("data/train.txt"), batch_size=self.args.train_batch_size,
                                           shuffle=True, drop_last=True, **kwargs)
        self.test_dataloader = DataLoader(ACEDataset("data/test.txt"), batch_size=self.args.test_batch_size,
                                          shuffle=True, drop_last=True, **kwargs)

        self.model = ResnetEncoderDecoder()
        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))
            cudnn.benchmark = True
            cudnn.enabled = True

        if self.args.resume:
            try:
                print("\nload the weight from pretrained-weight file")
                model_dict = self.model.state_dict()
                checkpoint = torch.load(self.args.pretrained_weight)['model_state_dict']
                model_dict.update(checkpoint)
                self.model.load_state_dict(model_dict, strict=True)
                print("Restoring the weight from pretrained-weight file \nFinished loading the weight\n")
            except Exception as e:
                raise e

        self.optimizer = optim.Adam(self.model.parameters(), lr=self.args.lr, weight_decay=5e-4)
        self.criterion = ACELoss()

    def work(self):
        torch.cuda.empty_cache()
        for epoch in range(1, self.args.epochs):
            self.train(epoch)
            self.test()
        torch.cuda.empty_cache()
        print("model finish training")

    def train(self, epoch):
        self.model.train()
        average_loss = []
        pbar = tqdm(self.train_dataloader, desc=f'Train Epoch:{epoch}/{self.args.epochs}')

        for data, target in pbar:
            data, target = data.to(self.device), target.long().to(self.device)
            self.optimizer.zero_grad()
            output = self.model(data)

            loss = self.criterion(output, target)
            loss.backward()
            self.optimizer.step()
            average_loss.append(loss.item())

            pbar.set_description(
                f'Train Epoch: {epoch}/{self.args.epochs} '
                f' train_loss: {np.mean(average_loss)} '
                f' learning_rate: {self.optimizer.state_dict()["param_groups"][0]["lr"]}'

            )
        global total_loss
        if np.mean(average_loss) < total_loss:
            torch.save({
                'model_state_dict': self.model.state_dict(),
            },
                './weights/best.pth')
            print("model saved")
            total_loss = np.mean(average_loss)

    @torch.no_grad()
    def inference(self):
        self.model.eval()

        image = Image.open("data/test/1.jpg").convert("1")
        image = image.resize([100, 100])
        image = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(0.5, 0.5)
        ])(image)

        image = image.unsqueeze(0)
        output = self.model(image)

        output = output[0]
        output = torch.softmax(output, -1)
        output = torch.argmax(output, -1).cpu()
        for out in output:
            mem_s = ""
            for o in out:
                mem_s += mem[int(o)]
            print(mem_s)


if __name__ == "__main__":
    model = Train()
    # model.work()
    model.inference()
